package edu.uaskl.jqf.model;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;

import edu.uaskl.jqf.algorithm.interfaces.Algorithm;
import edu.uaskl.jqf.algorithm.interfaces.Function;
import edu.uaskl.jqf.model.gates.interfaces.Gate;
import edu.uaskl.jqf.model.matrix.MatrixComplex;

/**
 * A quregister is an ordered set of qubits.<br>
 * <br>
 * Conventions typically used:<br>
 * |q1>, |q2>, |q3>, |q4> =<br>
 * |q1>|q2>|q3>|q4> =<br>
 * |q1q2q3q4><br>
 * <br>
 * Or with numbers:<br>
 * |1>, |1>, |0>, |1> =<br>
 * |1>|1>|0>|1> =<br>
 * |1101> =<br>
 * |13> =<br>
 * ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0)^t=<br>
 * [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,15] (indices) =><br>
 * (0,1)^t tensor (0,1)^t tensor (1,0)^t tensor (0,1)^t
 * 
 * @author tbach
 * 
 */
public class Quregister {
    private boolean qubitsAreValid = true; // this indicates if it is possible to work with the qubits list or not -tbach
    private final List<Qubit> qubits;
    private MatrixComplex qubitRegister; // TODO encapsulate?

    // TODO create constructor for string, e.g. qubit "001"

    /** Creates a {@link Quregister} with the given size and default state {@link Qubit#q0}. */
    public Quregister(final int size) {
        this.qubits = generateQubitsWithDefaultState0(size);
    }

    /** Creates a {@link Quregister} with the given qubits. */
    public Quregister(final Qubit... qubits) {
        this.qubits = new ArrayList<>(Arrays.asList(qubits));
    }

    /** Creates a {@link Quregister} with the given qubits. */
    public Quregister(final List<Qubit> qubits) {
        this.qubits = new ArrayList<>(qubits);
    }

    /**
     * Creates a {@link Quregister} with the given bits.<br>
     * Example:<br>
     * Input: [1,0,1]<br>
     * Result: Instance with |1>, |0>, |1>
     */
    public Quregister(final int... bits) {
        final List<Qubit> list = new ArrayList<>();
        for (final int bitsItem : bits)
            if (bitsItem == 0)
                list.add(Qubit.q0);
            else if (bitsItem == 1)
                list.add(Qubit.q1);
            else
                throw new IllegalArgumentException("0 or 1 expected, found: " + bitsItem);
        this.qubits = list;
    }

    /** Creates a {@link Quregister} with the given qubits. */
    public Quregister(final MatrixComplex qubitRegister) {
        this.qubitRegister = qubitRegister;
        this.qubitsAreValid = false;
        this.qubits = null;
    }

    /**
     * Returns a 2^n x 1 vector for n qubits, which is the tensor product of all qubits<br>
     * Example1:<br>
     * Qubits: |1>, |0> = (1,0)^t, (0,1)^t <br>
     * Result: (0,0,1,0)^t = |10><br>
     * <br>
     * Example2:<br>
     * Qubits: 1/sqrt(2) * (|0> + |1>), 1/sqrt(2) * (|0> - |1>) = 1/sqrt(2) * (1,1)^t, 1/sqrt(2) * (1,-1)^t <br>
     * Result: 1/2 * (1,-1,1,-1)^t = 1/2 * (|00> - |01> + |10> - |11>)<br>
     */
    public MatrixComplex getQubitRegister() {
        if (qubitRegister == null)
            createQubitRegister();
        return qubitRegister;
    }

    private synchronized void createQubitRegister() { // synchronized is not really needed, but we will do it to avoid large calculations -tbach
        if (qubitRegister != null) // prevent to recalculate everything -tbach
            return;
        qubitRegister = new MatrixComplex(1, 1, Complex.get10());
        for (final Qubit qubitsItem : qubits)
            qubitRegister = qubitRegister.tensor(qubitsItem);
    }

    private List<Qubit> generateQubitsWithDefaultState0(final int size) {
        final List<Qubit> result = new ArrayList<>(size);
        for (int i = 0; i < size; i++)
            result.add(Qubit.q0);
        return result;
    }

    public Quregister applyAlgorithm(final Algorithm algorithm) {
        /*
         * Well, we have multiple ways. 
         * 1) apply all gates to all qubits one after each other. This works only for single qubit operations 
         * 2) try to apply each gate after other, if we have gates 4x4, combine qubits with tensor and decompose
         * the result again - this is error prone and should be reviewed
         * 3) create kronecker product for all gates, apply to kronecker product for all qubits and find a way to
         * get and handle the resulting vector
         * 
         * --tbach
         */
        final int maxRows = getMaxColumnSizeForGatesFromAlgorithm(algorithm);
        if (maxRows == 2)
            return applyAlgorithmSingleQubits(algorithm);
        if (maxRows == 4)
            return applyAlgorithmAdvanced(algorithm);
        return applyAlgorithmFull(algorithm);
    }

    private int getMaxColumnSizeForGatesFromAlgorithm(final Algorithm algorithm) {
        int result = 0;
        for (final Gate gatesItem : algorithm.getGates())
            result = Math.max(result, gatesItem.getInstance().getRows());
        return result;
    }

    private Quregister applyAlgorithmSingleQubits(final Algorithm algorithm) {
        final List<Qubit> newQubits = new ArrayList<>(qubits.size());
        for (final Qubit qubitsItem : qubits) {
            Qubit newQubit = qubitsItem;
            for (final Gate gatesItem : algorithm.getGates())
                newQubit = gatesItem.getInstance().multiply(newQubit);
            newQubits.add(newQubit);
        }
        return new Quregister(newQubits);
    }

    private Quregister applyAlgorithmAdvanced(final Algorithm algorithm) { // TODO check if working? in general probably not -tbach
        System.err.println("algorithm advanced is not safe");
        List<Qubit> currentQubits = new ArrayList<>(qubits);
        int usedQubits = 0;
        for (final Gate gatesItem : algorithm.getGates()) {
            final MatrixComplex matrix = gatesItem.getInstance();
            final int neededSize = matrix.getColumns();
            if (neededSize == 2)
                currentQubits = applyMatrixForQubits(matrix, currentQubits);
            else if (neededSize == 4) {
                final List<Qubit> newQubits = new ArrayList<>(qubits.size());
                while (usedQubits < qubits.size()) {
                    if ((qubits.size() - usedQubits) <= 1)
                        throw new IllegalStateException("Not enough qubits left.");
                    final MatrixComplex tensorQubit = currentQubits.get(usedQubits++).tensor(currentQubits.get(usedQubits++));
                    final MatrixComplex result = matrix.multiply(tensorQubit);

                    // this is not unique, because kronecker product is not injective
                    final Complex v0 = result.get(0, 0);
                    final Complex v1 = result.get(1, 0);
                    final Complex v2 = result.get(2, 0);
                    final Complex v3 = result.get(3, 0);
                    final Complex a0 = v0.multiply(v0).add(v1.multiply(v1)).getSqrtPositive();
                    final Complex a1 = v2.multiply(v2).add(v3.multiply(v3)).getSqrtPositive();
                    final Complex b0 = v0.multiply(v0).add(v2.multiply(v2)).getSqrtPositive();
                    final Complex b1 = v1.multiply(v1).add(v3.multiply(v3)).getSqrtPositive();
                    Qubit newQubitFirst = new Qubit(a0, a1);
                    Qubit newQubitSecond = new Qubit(b0, b1);

                    for (int i = 0; i < 16; ++i) { // this is still not unique, but at least to a factor of (-1)
                        newQubitFirst = new Qubit(a0.multiply(1 - (2 * ((i >> 1) % 2))), a1.multiply(1 - (2 * ((i >> 0) % 2))));
                        newQubitSecond = new Qubit(b0.multiply(1 - (2 * ((i >> 3) % 2))), b1.multiply(1 - (2 * ((i >> 2) % 2))));
                        if (newQubitFirst.tensor(newQubitSecond).equals(result))
                            System.out.println(newQubitFirst + " # " + newQubitSecond);
                    }

                    newQubits.add(newQubitFirst);
                    newQubits.add(newQubitSecond);
                }
                currentQubits = newQubits;
            } else
                throw new IllegalStateException("needed size != (2 or 4). neededSize: " + neededSize);
        }
        return new Quregister(currentQubits);
    }

    private List<Qubit> applyMatrixForQubits(final MatrixComplex matrix, final List<Qubit> qubitsSource) {
        final List<Qubit> newQubits = new ArrayList<>(qubitsSource.size());
        for (final Qubit qubitsItem : qubitsSource)
            newQubits.add(matrix.multiply(qubitsItem));
        return newQubits;
    }

    /**
     * This method multiplies the full gate size from the algorithm with the full qubit space representation.
     */
    private Quregister applyAlgorithmFull(final Algorithm algorithm) {
        MatrixComplex currentQubitRegister = getQubitRegister();
        if (currentQubitRegister.getColumns() != 1)
            throw new IllegalStateException("Register columns != 1. Dimension: " + currentQubitRegister.getDimensionAsString());
        final int size = currentQubitRegister.getRows();
        for (final Gate gatesItem : algorithm.getGates())
            currentQubitRegister = gatesItem.getSize(size).multiply(currentQubitRegister);
        return new Quregister(currentQubitRegister);
    }

    public Quregister applyFunction(final Function function) {
        return function.apply(this);
    }

    public List<Qubit> getQubits() {
        return Collections.unmodifiableList(qubits);
    }

    @Override
    public String toString() {
        if (qubitsAreValid) {
            final StringBuilder stringBuilder = new StringBuilder();
            // TODO does not work for invalid qubits
            for (final Qubit qubitsItem : qubits)
                stringBuilder.append(qubitsItem).append(" # ");
            return stringBuilder.toString();
        }
        return qubitRegister.toString();
    }

    public String measure() {
        if (qubitsAreValid)
            return measureQubits();
        return measureQubitRegister();
    }

    private String measureQubits() {
        final StringBuilder stringBuilder = new StringBuilder();
        for (final Qubit qubitsItem : qubits)
            stringBuilder.append(qubitsItem.measure());
        return stringBuilder.toString();
    }

    private String measureQubitRegister() { // exactly the same as the other implementation -tbach
        /* plan: we have a random threshold between 0 and 1
         * we go over all rows of the result vector, add them up until we reach the random threshold
         * this is the random selected result state
         * -tbach
         */
        final MatrixComplex currentQubitRegister = getQubitRegister();
        final double threshold = ThreadLocalRandom.current().nextDouble();
        double sum = 0;
        for (int i = 0; i < currentQubitRegister.getRows(); ++i) {
            // well, we could have some problems with accumulated fp addition errors. but ok, we will not care about it now -tbach
            sum += currentQubitRegister.get(i, 0).getAbsSquare();
            if (sum >= threshold)
                return Integer.toBinaryString(i);
        }
        throw new IllegalStateException("Probability distribution failed. Threshould: " + threshold + ", sum: " + sum + ", qubitRegister: "
                + currentQubitRegister);
    }

    public Integer measureAsDecimal() {
        return Integer.valueOf(measure(), 2);
    }
}
